import backtrader as bt
import numpy as np
import pandas as pd
from typing import Dict, Any, Optional

class VolatilityStrategy(bt.Strategy):
    """
    波动率策略
    
    策略逻辑：
    1. 使用ATR计算真实波动率
    2. 在低波动率期间买入
    3. 在高波动率期间卖出
    4. 基于波动率调整仓位大小
    5. 设置动态止损
    """
    
    params = (
        ('atr_period', 14),         # ATR周期
        ('volatility_lookback', 20), # 波动率回看周期
        ('low_vol_threshold', 0.8),  # 低波动率阈值（相对于历史平均）
        ('high_vol_threshold', 1.5), # 高波动率阈值
        ('position_size', 0.1),      # 基础仓位比例
        ('max_position_size', 0.25), # 最大仓位比例
        ('stop_loss_atr', 2.0),      # 基于ATR的止损倍数
    )
    
    def __init__(self):
        # 波动率指标
        self.atr = bt.indicators.ATR(
            self.data,
            period=self.params.atr_period
        )
        
        # 计算历史平均ATR
        self.atr_ma = bt.indicators.SMA(
            self.atr,
            period=self.params.volatility_lookback
        )
        
        # 价格均线
        self.price_ma = bt.indicators.SMA(
            self.data.close,
            period=20
        )
        
        # 交易状态
        self.order = None
        self.stop_order = None
        self.entry_price = None
        
    def next(self):
        if self.order:
            return
            
        if not self.position:
            # 寻找买入机会
            current_atr = self.atr[0]
            avg_atr = self.atr_ma[0]
            current_price = self.data.close[0]
            price_ma = self.price_ma[0]
            
            # 低波动率且价格在均线附近
            volatility_ratio = current_atr / avg_atr
            price_near_ma = abs(current_price - price_ma) / price_ma < 0.05
            
            if (volatility_ratio < self.params.low_vol_threshold and 
                price_near_ma and current_price > price_ma):
                
                # 根据波动率调整仓位
                vol_adjustment = 1 - volatility_ratio
                adjusted_position_size = min(
                    self.params.position_size * (1 + vol_adjustment),
                    self.params.max_position_size
                )
                
                # 计算买入数量
                cash = self.broker.getcash()
                size = int((cash * adjusted_position_size) / current_price)
                
                if size > 0:
                    self.order = self.buy(size=size)
                    self.entry_price = current_price
                    print(f'低波动率买入: 价格={current_price:.2f}, ATR比率={volatility_ratio:.3f}')
        
        else:
            # 寻找卖出机会
            current_atr = self.atr[0]
            avg_atr = self.atr_ma[0]
            volatility_ratio = current_atr / avg_atr
            
            # 高波动率或价格跌破均线
            current_price = self.data.close[0]
            price_ma = self.price_ma[0]
            
            if (volatility_ratio > self.params.high_vol_threshold or 
                current_price < price_ma * 0.98):
                
                self.order = self.sell(size=self.position.size)
                print(f'高波动率卖出: 价格={current_price:.2f}, ATR比率={volatility_ratio:.3f}')
    
    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return
            
        if order.status in [order.Completed]:
            if order.isbuy():
                # 设置基于ATR的动态止损
                price = order.executed.price
                atr_value = self.atr[0]
                stop_price = price - (atr_value * self.params.stop_loss_atr)
                
                self.stop_order = self.sell(
                    size=order.executed.size,
                    exectype=bt.Order.Stop,
                    price=stop_price
                )
                
                print(f'买入完成: 价格={price:.2f}, 数量={order.executed.size}, 止损={stop_price:.2f}')
            else:
                # 取消止损
                if self.stop_order:
                    self.cancel(self.stop_order)
                    
                print(f'卖出完成: 价格={order.executed.price:.2f}, 数量={order.executed.size}')
        
        self.order = None

def get_strategy_params() -> Dict[str, Any]:
    """获取策略参数"""
    return {
        'name': '波动率策略',
        'description': '基于ATR波动率的低买高卖策略，在低波动率时买入，高波动率时卖出',
        'parameters': {
            'atr_period': {'type': 'int', 'default': 14, 'min': 10, 'max': 21, 'description': 'ATR周期'},
            'volatility_lookback': {'type': 'int', 'default': 20, 'min': 15, 'max': 30, 'description': '波动率回看周期'},
            'low_vol_threshold': {'type': 'float', 'default': 0.8, 'min': 0.6, 'max': 1.0, 'description': '低波动率阈值'},
            'high_vol_threshold': {'type': 'float', 'default': 1.5, 'min': 1.2, 'max': 2.0, 'description': '高波动率阈值'},
            'position_size': {'type': 'float', 'default': 0.1, 'min': 0.05, 'max': 0.2, 'description': '基础仓位比例'},
            'max_position_size': {'type': 'float', 'default': 0.25, 'min': 0.15, 'max': 0.4, 'description': '最大仓位比例'},
            'stop_loss_atr': {'type': 'float', 'default': 2.0, 'min': 1.5, 'max': 3.0, 'description': '基于ATR的止损倍数'}
        }
    } 